package com.yilin.gateway.filter_factory;

import com.google.gson.reflect.TypeToken;
import com.yilin.gateway.core.BaseRuntimeException;
import com.yilin.gateway.core.EncryptUtils;
import com.yilin.gateway.core.JsonUtil;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.factory.rewrite.CachedBodyOutputMessage;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpHeaders;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;

public class DecryptRequestBodyGatewayFilterFactory extends AbstractGatewayFilterFactory implements Ordered {
    private Logger logger = LoggerFactory.getLogger(DecryptRequestBodyGatewayFilterFactory.class);
    private final static String TOKEN = "token";
    private final static String UPLOAD = "upload";
    private final static String ContentType = "Content-Type";

    private final List<HttpMessageReader<?>> messageReaders;

    public DecryptRequestBodyGatewayFilterFactory() {
        this.messageReaders = HandlerStrategies.withDefaults().messageReaders();
    }

    @Override
    @SuppressWarnings("unchecked")
    public GatewayFilter apply(Object config) {
        return (exchange, chain) -> {
            ServerRequest serverRequest = ServerRequest.create(exchange, this.messageReaders);
            String tokenID = exchange.getRequest().getHeaders().getFirst(TOKEN);
            String upload = exchange.getRequest().getHeaders().getFirst(UPLOAD);
            String contentType = exchange.getRequest().getHeaders().getFirst(ContentType);
            // ### 上传文件时不要重写BODY,防止文件损坏
            if (Boolean.parseBoolean(upload) || (contentType!=null&&contentType.contains("multipart/form-data"))) return chain.filter(exchange);
            Mono<String> modifiedBody = serverRequest.bodyToMono(String.class).flatMap(originalBody -> modifyBody(tokenID).apply(exchange, Mono.just(originalBody)));
            BodyInserter bodyInserter = BodyInserters.fromPublisher(modifiedBody, String.class);
            HttpHeaders headers = new HttpHeaders();
            headers.putAll(exchange.getRequest().getHeaders());
            headers.remove(HttpHeaders.CONTENT_LENGTH);
            CachedBodyOutputMessage outputMessage = new CachedBodyOutputMessage(exchange, headers);
            return bodyInserter.insert(outputMessage, new BodyInserterContext()).then(Mono.defer(() -> {
                ServerHttpRequest decorator = decorate(exchange, headers, outputMessage);
                return chain.filter(exchange.mutate().request(decorator).build());
            }));
        };
    }

    private BiFunction<ServerWebExchange, Mono<String>, Mono<String>> modifyBody(String token) {
        return (exchange, body) -> {
            try {
                AtomicReference<String> result = new AtomicReference<>();
                body.subscribe(value -> {
                    if (value!=null && JsonUtil.isJson(value)){
                        Map<String,Object> map = JsonUtil.jsonToObject(value,new TypeToken<Map<String,Object>>(){});
                        if (map.containsKey("body")){
                            String o_data = (String) map.get("body");
                            value = EncryptUtils.aesDecrypt(o_data, token);
                        }
                    }
                    result.set(value);
                }, e -> logger.error(e.getMessage(), e));
                return Mono.just(result.get());
            } catch (Exception e) {
                logger.error("gateway parameter decryption exception",e);
                throw new BaseRuntimeException("Parameter decryption exception");
            }
        };
    }


    private ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers, CachedBodyOutputMessage outputMessage) {
        return new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public HttpHeaders getHeaders() {
                long contentLength = headers.getContentLength();
                HttpHeaders httpHeaders = new HttpHeaders();
                httpHeaders.putAll(super.getHeaders());
                if (contentLength > 0) {
                    httpHeaders.setContentLength(contentLength);
                } else {
                    httpHeaders.set(HttpHeaders.TRANSFER_ENCODING, "chunked");
                }
                return httpHeaders;
            }
            @Override
            public Flux<DataBuffer> getBody() {
                return outputMessage.getBody();
            }
        };

    }

    @Override
    public int getOrder() {
        return -2;
    }
}